import cv2
import os
import numpy as np
from concurrent.futures import ThreadPoolExecutor
from queue import Queue

def detect_frame_changes(video_path, output_dir, threshold=30, scale_factor=0.5, batch_size=10):


    cuda_available = cv2.cuda.getCudaEnabledDeviceCount() > 0

    if cuda_available:
        print("CUDA可用")
    else:
        print("CUDA不可用，切换到CPU")  

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 初始化GPU加速
    cv2.setUseOptimized(True)
    cv2.setNumThreads(4)
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print("无法打开视频文件")
        return

    # 获取视频信息
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) * scale_factor)
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * scale_factor)
    
    # 初始化变量
    frame_queue = Queue(maxsize=batch_size * 2)
    save_queue = Queue()
    executor = ThreadPoolExecutor(max_workers=16)
    
    def process_frame(frame, prev_frame):
        # 下采样
        frame = cv2.resize(frame, (width, height))
        prev_frame = cv2.resize(prev_frame, (width, height))
        # 转换为灰度图像
        curr_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)

        
        # 检查CUDA是否可用
        cuda_available = cv2.cuda.getCudaEnabledDeviceCount() > 0
        
        if cuda_available:
            #print("CUDA可用")
            frame_diff = cv2.cuda.absdiff(curr_gray, prev_gray)
            return frame_diff.mean()
        else:
            #print("CUDA不可用，切换到CPU")
            # 计算帧间差异
            frame_diff = cv2.absdiff(curr_gray, prev_gray)
            return frame_diff.mean()

    def save_frames(frames):
        for frame, path in frames:
            cv2.imwrite(path, frame)
    
    # 读取第一帧
    ret, prev_frame = cap.read()
    if not ret:
        print("无法读取视频帧")
        return

    frame_count = 0
    saved_count = 0
    batch = []
    i = 0
    while True:
        i += 1
        
     
        ret, curr_frame = cap.read()
        if (i % 10 != 0):
            continue
        if not ret:
            break
            
        # 提交任务到线程池
        future = executor.submit(process_frame, curr_frame, prev_frame)
        diff_score = future.result()
        
        if diff_score > threshold:
            output_path = f"{output_dir}/frame_{frame_count:04d}.jpg"
            batch.append((curr_frame, output_path))
            print(f"保存帧 {frame_count:04d}，差异值: {diff_score:.2f}, 保存到 {output_path}")
            cv2.imwrite(output_path, curr_frame)
            #if len(batch) >= batch_size:
            # save_queue.put(batch.copy())
            # executor.submit(save_frames, batch)
            # batch.clear()
            
            saved_count += 1
            
        # 更新前一帧
        prev_frame = curr_frame
        frame_count += 1
        
        # 显示进度
        if frame_count % 100 == 0:
            current_seconds = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000
            total_seconds = total_frames / fps
            print(f"进度: {current_seconds:.1f}/{total_seconds:.1f} 秒 ({current_seconds/total_seconds:.1%})")
    
    # # 保存剩余帧
    # if batch:
    #     save_queue.put(batch)
    #     executor.submit(save_frames, batch)
    
    # 释放资源
    cap.release()
    executor.shutdown()
    print(f"分析完成，共保存 {saved_count} 帧。")

# 使用示例
video_path = "inputs/第三节：大模型的应用.mp4"  # 输入视频文件路径
output_dir = "output_frames"  # 输出目录
detect_frame_changes(video_path, output_dir, scale_factor=0.25, threshold=10)